from ultralytics import YOLO
import os
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description='Train YOLOv8 Fruit Recognition Model')
    parser.add_argument('--data', type=str, default='data/data_fixed.yaml', help='Data config file path')
    parser.add_argument('--epochs', type=int, default=100, help='Number of epochs')
    parser.add_argument('--batch', type=int, default=16, help='Batch size')
    parser.add_argument('--imgsz', type=int, default=640, help='Image size')
    parser.add_argument('--model', type=str, default='yolov8n.pt', help='Model type')
    parser.add_argument('--device', type=str, default='', help='Training device, e.g. 0 or 0,1,2,3 or cpu')
    return parser.parse_args()

def main():
    args = parse_args()

    # Ensure model directory exists
    os.makedirs('models', exist_ok=True)
    
    # Load pretrained model
    model = YOLO(args.model)

    # Train model
    results = model.train(
        data=args.data,
        epochs=args.epochs,
        batch=args.batch,
        imgsz=args.imgsz,
        device=args.device,
        project='runs',
        name='train',
        exist_ok=True
    )

    # Validate model
    model.val()

    # Copy best model to models directory
    best_model_path = os.path.join('runs', 'train', 'weights', 'best.pt')
    if os.path.exists(best_model_path):
        import shutil
        shutil.copy(best_model_path, 'models/best.pt')
        print(f'Best model copied to models/best.pt')

    print('Training completed!')

if __name__ == '__main__':
    main()
